import cv2
import numpy as np
import pandas as pd
import os
from scipy.fftpack import fft2

# 数据预处理
def preprocess_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.resize(image, (256, 256))  # 归一化处理
    return image

# 颜色偏移分析
def color_shift_analysis(image):
    mean_r = np.mean(image[:, :, 0])
    mean_g = np.mean(image[:, :, 1])
    mean_b = np.mean(image[:, :, 2])
    d_rgb = max(abs(mean_r - mean_g), abs(mean_g - mean_b), abs(mean_b - mean_r))
    return d_rgb

# 低光照分析
def low_light_analysis(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    l_mean = np.mean(gray)
    l_std = np.std(gray)
    return l_mean, l_std

# 模糊分析
def blur_analysis(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    laplacian_var = cv2.Laplacian(gray, cv2.CV_64F).var()
    fft_energy = np.sum(np.abs(fft2(gray))**2)
    return laplacian_var, fft_energy

# 主程序
def main():
    image_dir = r"E:\Attachment\Attachment 1"
    output_path = r"C:\Users\86180\Desktop\Answer\1\Answer.xlsx"
    
    results = []
    total_images = 0
    color_shift_count = 0
    low_light_count = 0
    blur_count = 0
    clear_count = 0
    
    for filename in os.listdir(image_dir):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            total_images += 1
            image_path = os.path.join(image_dir, filename)
            image = preprocess_image(image_path)
            
            d_rgb = color_shift_analysis(image)
            l_mean, l_std = low_light_analysis(image)
            laplacian_var, fft_energy = blur_analysis(image)
            
            # 动态计算优先级
            scores = {
                "Color Shift": d_rgb / some_threshold,
                "Low Light": (some_light_threshold - l_mean) / some_light_threshold,
                "Blur": min(laplacian_var / some_blur_threshold, fft_energy / some_fft_threshold)
            }
            
            # 选择得分最高的分类
            classification = max(scores, key=scores.get)
            if scores[classification] <= 0:
                classification = "Clear"
                clear_count += 1
            else:
                if classification == "Color Shift":
                    color_shift_count += 1
                elif classification == "Low Light":
                    low_light_count += 1
                elif classification == "Blur":
                    blur_count += 1
            
            results.append({
                "image file name": filename,
                "Degraded Image Classification": classification
            })
    
    df = pd.DataFrame(results)
    df.to_excel(output_path, index=False)

    # 打印统计信息
    print(f"Total images: {total_images}")
    print(f"Color Shift: {color_shift_count}")
    print(f"Low Light: {low_light_count}")
    print(f"Blur: {blur_count}")
    print(f"Clear: {clear_count}")

if __name__ == "__main__":
    # 阈值设置
    some_threshold = 30  # 颜色偏移阈值
    some_light_threshold = 50  # 低光照阈值
    some_blur_threshold = 100  # 拉普拉斯方差阈值
    some_fft_threshold = 1e6  # FFT能量阈值
    
    main()
